package com.qiandao.xuexitong.websocket;

import com.alibaba.fastjson.JSON;

import com.qiandao.xuexitong.constant.JwtClaimsConstant;
import com.qiandao.xuexitong.exception.BusinessException;
import com.qiandao.xuexitong.properties.JwtProperties;
import com.qiandao.xuexitong.util.JwtUtil;
import io.jsonwebtoken.Claims;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.security.KeyStore;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
@ServerEndpoint("/ws/notification")
public class WebSocketServer {
    
    private static JwtProperties jwtProperties;
    
    @Autowired
    public void setJwtProperties(JwtProperties jwtProperties) {
        WebSocketServer.jwtProperties = jwtProperties;
    }

    /**
     * 存储所有连接的客户端
     */
    private static final Map<Long, Session> sessionMap = new ConcurrentHashMap<>();
    
    /**
     * 连接建立时调用
     */
    @OnOpen
    public void onOpen(Session session) {
        try {
            // 从请求参数中获取token
            Map<String, List<String>> params = session.getRequestParameterMap();
            if (params == null || !params.containsKey("token")) {
                log.error("WebSocket连接失败：未提供token");
                session.close();
                return;
            }

            List<String> tokens = params.get("token");
            if (tokens == null || tokens.isEmpty()) {
                log.error("WebSocket连接失败：token为空");
                session.close();
                return;
            }

            String token = tokens.get(0);
            if (token == null || token.isEmpty()) {
                log.error("WebSocket连接失败：token为空");
                session.close();
                return;
            }

            // 验证token并获取用户ID
            Claims claims = JwtUtil.parseJWT(jwtProperties.getSecretKey(), token);
            Long userId = Long.valueOf(claims.get(JwtClaimsConstant.USER_ID).toString());
            if (userId == null) {
                log.error("WebSocket连接失败：无效的token");
                session.close();
                return;
            }

            // 存储会话
            sessionMap.put(userId, session);
            log.info("WebSocket连接建立，用户ID: {}", userId);
        } catch (Exception e) {
            log.error("WebSocket连接建立失败: {}", e.getMessage(), e);
            try {
                session.close();
            } catch (IOException ex) {
                log.error("关闭WebSocket连接失败", ex);
            }
        }
    }
    
    /**
     * 收到客户端消息时调用
     */
    @OnMessage
    public void onMessage(String message, Session session) {
        log.info("收到客户端消息: {}", message);
        // 处理心跳消息
        if ("ping".equals(message)) {
            try {
                session.getBasicRemote().sendText("pong");
            } catch (IOException e) {
                log.error("发送心跳响应失败", e);
            }
        }
    }
    
    /**
     * 连接关闭时调用
     */
    @OnClose
    public void onClose(Session session) {
        // 从sessionMap中移除该会话
        sessionMap.entrySet().removeIf(entry -> entry.getValue().equals(session));
        log.info("WebSocket连接关闭");
    }
    
    /**
     * 发生错误时调用
     */
    @OnError
    public void onError(Session session, Throwable error) {
        log.error("WebSocket发生错误: {}", error.getMessage(), error);
        try {
            session.close();
        } catch (IOException e) {
            log.error("关闭WebSocket连接失败", e);
        }
    }
    
    /**
     * 向多个用户发送消息
     * @param userIds 用户ID列表
     * @param message 消息内容
     */
    public static void sendMessageToUsers(List<Long> userIds, Object message) {
        try {
            String jsonMessage = JSON.toJSONString(message);
            for (Long userId : userIds) {
                Session session = sessionMap.get(userId);
                if (session != null && session.isOpen()) {
                    session.getBasicRemote().sendText(jsonMessage);
                }
            }
        } catch (Exception e) {
            log.error("发送WebSocket消息失败: {}", e.getMessage(), e);
        }
    }
} 